{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dfea1fe5",
   "metadata": {},
   "source": [
    "# jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d2ce6c58",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import tir\n",
    "from tvm.relax.frontend.nn import spec\n",
    "from tvm.relax.frontend import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fccac699",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Layer(nn.Module):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def forward(self, x: nn.Tensor):\n",
    "        y = nn.add(x, x)\n",
    "        return y\n",
    "\n",
    "forward_spec = {\"forward\": {\"x\": spec.Tensor([10, 5], dtype=\"float32\")}}\n",
    "mod = Layer()\n",
    "\n",
    "for debug in [False, True]:\n",
    "    model = mod.jit(spec=forward_spec, debug=debug)\n",
    "\n",
    "    x = torch.rand((10, 5), dtype=torch.float32)\n",
    "    y = model[\"forward\"](x)\n",
    "    assert isinstance(y, torch.Tensor)\n",
    "    assert torch.allclose(x + x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4654a0e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Layer(nn.Module):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def forward(self, x: nn.Tensor, i: tir.Var):\n",
    "        y = nn.add(x, x)\n",
    "        y = nn.reshape(y, (i, 5, 5))\n",
    "        return y\n",
    "\n",
    "forward_spec = {\"forward\": {\"x\": spec.Tensor([10, 5], dtype=\"float32\"), \"i\": int}}\n",
    "mod = Layer()\n",
    "\n",
    "for debug in [False, True]:\n",
    "    model = mod.jit(spec=forward_spec, debug=debug)\n",
    "\n",
    "    x = torch.rand((10, 5), dtype=torch.float32)\n",
    "    y = model[\"forward\"](x, 2)\n",
    "    assert isinstance(y, torch.Tensor)\n",
    "    assert torch.allclose(torch.reshape(x + x, (2, 5, 5)), y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c8ba65f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Layer(nn.Module):\n",
    "    def __init__(self):\n",
    "        self.cache = nn.KVCache(10, [10, 5])\n",
    "\n",
    "    def forward(self, x: nn.Tensor, total_seq_len: tir.Var):\n",
    "        self.cache.append(x)\n",
    "        y = self.cache.view(total_seq_len)\n",
    "        return y\n",
    "\n",
    "forward_spec = {\n",
    "    \"forward\": {\"x\": spec.Tensor([1, 10, 5], dtype=\"float32\"), \"total_seq_len\": int}\n",
    "}\n",
    "mod = Layer()\n",
    "\n",
    "for debug in [False, True]:\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        model = mod.jit(spec=forward_spec, debug=debug)\n",
    "\n",
    "    x0 = torch.rand((1, 10, 5), dtype=torch.float32)\n",
    "    y = model[\"forward\"](x0, 1)\n",
    "    assert isinstance(y, torch.Tensor)\n",
    "    assert torch.allclose(x0, y)\n",
    "\n",
    "    x1 = torch.rand((1, 10, 5), dtype=torch.float32)\n",
    "    y = model[\"forward\"](x1, 2)\n",
    "    assert torch.allclose(torch.concat([x0, x1], dim=0), y)\n",
    "\n",
    "    x2 = torch.rand((1, 10, 5), dtype=torch.float32)\n",
    "    y = model[\"forward\"](x2, 3)\n",
    "    assert torch.allclose(torch.concat([x0, x1, x2], dim=0), y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "77313a00",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Layer(nn.Module):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def forward(self, x: tuple[nn.Tensor, nn.Tensor]):\n",
    "        assert isinstance(x, tuple)\n",
    "        x0 = x[0]\n",
    "        x1 = x[1]\n",
    "        y0 = nn.add(x0, x1)\n",
    "        y1 = nn.subtract(x0, x1)\n",
    "        return (y0, y1)\n",
    "\n",
    "forward_spec = {\n",
    "    \"forward\": {\n",
    "        \"x\": (\n",
    "            spec.Tensor([10, 5], dtype=\"float32\"),\n",
    "            spec.Tensor([10, 5], dtype=\"float32\"),\n",
    "        )\n",
    "    }\n",
    "}\n",
    "mod = Layer()\n",
    "\n",
    "for debug in [False, True]:\n",
    "    model = mod.jit(spec=forward_spec, debug=debug)\n",
    "\n",
    "    x0 = torch.rand((10, 5), dtype=torch.float32)\n",
    "    x1 = torch.rand((10, 5), dtype=torch.float32)\n",
    "    x = (x0, x1)\n",
    "    y = model[\"forward\"](x)\n",
    "\n",
    "    assert torch.allclose(x0 + x1, y[0])\n",
    "    assert torch.allclose(x0 - x1, y[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6b5c317d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Layer(nn.Module):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def forward(self, x: list[nn.Tensor]):\n",
    "        assert isinstance(x, list)\n",
    "        x0 = x[0]\n",
    "        x1 = x[1]\n",
    "        y0 = nn.add(x0, x1)\n",
    "        y1 = nn.subtract(x0, x1)\n",
    "        return (y0, y1)\n",
    "\n",
    "forward_spec = {\n",
    "    \"forward\": {\n",
    "        \"x\": [\n",
    "            spec.Tensor([10, 5], dtype=\"float32\"),\n",
    "            spec.Tensor([10, 5], dtype=\"float32\"),\n",
    "        ]\n",
    "    }\n",
    "}\n",
    "mod = Layer()\n",
    "\n",
    "for debug in [False, True]:\n",
    "    model = mod.jit(spec=forward_spec, debug=debug)\n",
    "\n",
    "    x0 = torch.rand((10, 5), dtype=torch.float32)\n",
    "    x1 = torch.rand((10, 5), dtype=torch.float32)\n",
    "    x = (x0, x1)\n",
    "    y = model[\"forward\"](x)\n",
    "\n",
    "    assert torch.allclose(x0 + x1, y[0])\n",
    "    assert torch.allclose(x0 - x1, y[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9b9a82c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "class Layer(nn.Module):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def forward(self, x: Tuple[nn.Tensor, nn.Tensor, int]):\n",
    "        x0 = x[0]\n",
    "        x1 = x[1]\n",
    "        y0 = nn.add(x0, x1)\n",
    "        y1 = nn.subtract(x0, x1)\n",
    "        y2 = nn.reshape(x0, (5, x[2], 5))\n",
    "        return (y0, y1, y2)\n",
    "\n",
    "forward_spec = {\n",
    "    \"forward\": {\n",
    "        \"x\": (spec.Tensor([10, 5], dtype=\"float32\"), spec.Tensor([10, 5], dtype=\"float32\"), int)\n",
    "    }\n",
    "}\n",
    "mod = Layer()\n",
    "\n",
    "for debug in [False, True]:\n",
    "    model = mod.jit(spec=forward_spec, debug=debug)\n",
    "\n",
    "    x0 = torch.rand((10, 5), dtype=torch.float32)\n",
    "    x1 = torch.rand((10, 5), dtype=torch.float32)\n",
    "    x = (x0, x1, 2)\n",
    "    y0, y1, y2 = model[\"forward\"](x)\n",
    "\n",
    "    assert torch.allclose(x0 + x1, y0)\n",
    "    assert torch.allclose(x0 - x1, y1)\n",
    "    assert torch.allclose(torch.reshape(x0, (5, 2, 5)), y2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "188009e4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
